# markov.py
# Johannes Kaisinger, 4 July 2024
#
# Demonstrate estimation of markov parameters.
# SISO, SIMO, MISO, MIMO case

# 导入必要的库
import numpy as np
import matplotlib.pyplot as plt
import os

import control as ct

def create_impulse_response(H, time, transpose, dt):
    """
    辅助函数，用于使用 TimeResponseData 类型进行绘图。

    参数:
    H (array-like): 脉冲响应矩阵
    time (array-like): 时间向量
    transpose (bool): 是否转置脉冲响应矩阵
    dt (float): 采样时间间隔

    返回:
    ct.TimeResponseData: 包含脉冲响应信息的 TimeResponseData 对象
    """
    # 将输入的 H 转换为至少三维的 numpy 数组
    H = np.array(H, ndmin=3)

    # 如果需要，转置脉冲响应矩阵
    if transpose:
        H = np.transpose(H)
    
    # 获取脉冲响应矩阵的维度
    q, p, m = H.shape
    # 初始化输入矩阵
    inputs = np.zeros((p,p,m))

    # 判断系统是否为单输入单输出 (SISO) 系统
    issiso = True if (q == 1 and p == 1) else False
    
    # 初始化输入标签列表
    input_labels = []
    # 初始化轨迹标签和轨迹类型列表
    trace_labels, trace_types = [], []
    for i in range(p):
        # 设置单位面积脉冲输入
        inputs[i,i,0] = 1/dt 
        input_labels.append(f"u{[i]}")
        trace_labels.append(f"From u{[i]}")
        trace_types.append('impulse')

    # 初始化输出标签列表
    output_labels = []
    for i in range(q):
        output_labels.append(f"y{[i]}")

    # 创建并返回 TimeResponseData 对象
    return ct.TimeResponseData(time=time[:m],
                            outputs=H,
                            output_labels=output_labels,
                            inputs=inputs,
                            input_labels=input_labels,
                            trace_labels=trace_labels,
                            trace_types=trace_types,
                            sysname="H_est",
                            transpose=transpose,
                            plot_inputs=False,
                            issiso=issiso)

# 设置一个质量 - 弹簧 - 阻尼器系统 (2 自由度，多输入多输出情况)
# 参考书籍: Mechanical Vibrations: Theory and Application, SI Edition, 1st ed.
# 图 6.5 / 示例 6.7
# 动力学方程: m q_dd + c q_d + k q = f
m1, k1, c1 = 1., 4., 1.
m2, k2, c2 = 2., 2., 1.
k3, c3 = 6., 2.

# 定义系统状态矩阵 A
A = np.array([
    [0., 0., 1., 0.],
    [0., 0., 0., 1.],
    [-(k1+k2)/m1, (k2)/m1, -(c1+c2)/m1, c2/m1],
    [(k2)/m2, -(k2+k3)/m2, c2/m2, -(c2+c3)/m2]
])
# 定义系统输入矩阵 B
B = np.array([[0.,0.],[0.,0.],[1/m1,0.],[0.,1/m2]])
# 定义系统输出矩阵 C
C = np.array([[1.0, 0.0, 0.0, 0.0],[0.0, 1.0, 0.0, 0.0]])
# 定义系统前馈矩阵 D
D = np.zeros((2,2))

# 定义系统类型列表
xixo_list = ["SISO","SIMO","MISO","MIMO"]
# 选择用于估计的系统类型
xixo = xixo_list[3] 

# 根据选择的系统类型创建状态空间模型
match xixo:
    case "SISO":
        sys = ct.StateSpace(A, B[:,0], C[0,:], D[0,0])
    case "SIMO":
        sys = ct.StateSpace(A, B[:,:1], C, D[:,:1])
    case "MISO":
        sys = ct.StateSpace(A, B, C[:1,:], D[:1,:])
    case "MIMO":
        sys = ct.StateSpace(A, B, C, D)

# 定义采样时间间隔
dt = 0.25
# 将连续系统离散化
sysd = sys.sample(dt, method='zoh')
# 设置离散系统的名称
sysd.name = "H_true"

# 生成随机输入信号
t = np.arange(0,100,dt)
u = np.random.randn(sysd.B.shape[-1], len(t))

# 计算系统的强迫响应
response = ct.forced_response(sysd, U=u)
# 绘制强迫响应曲线
response.plot()
plt.show()

# 定义脉冲响应的长度
m = 50
# 计算系统的真实脉冲响应
ir_true = ct.impulse_response(sysd, T=dt*m)

# 估计马尔可夫参数
H_est = ct.markov(response, m, dt=dt)
# 辅助函数，仅用于绘图
ir_est = create_impulse_response(H_est,
                                 ir_true.time,
                                 ir_true.transpose,
                                 dt)

# 绘制真实脉冲响应曲线
ir_true.plot(title=xixo)
# 绘制估计的脉冲响应曲线
ir_est.plot(color='orange',linestyle='dashed')
plt.show()

# 如果环境变量中不包含 'PYCONTROL_TEST_EXAMPLES'，则显示图形
if 'PYCONTROL_TEST_EXAMPLES' not in os.environ:
    plt.show()